Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Make PyBaMM importable in Linux systems if jax is already installed #1874

Merged
merged 10 commits into from
Dec 29, 2021
Merged

Make PyBaMM importable in Linux systems if jax is already installed #1874

merged 10 commits into from
Dec 29, 2021

Conversation

Saransh-cpp
Copy link
Member

@Saransh-cpp Saransh-cpp commented Dec 27, 2021 •
edited
Loading

Description

  • Check for the version of jax in have_jax()
  • Give an error in install_jax() if jax is already available with a different version

Additionally, the cell with pybamm.print_citations() gives an error which is described here - pybamm-team/BattBot#47 - and is most probably an error within the pybtex library.

Fixes # (issue)

Type of change

Please add a line in the relevant section of CHANGELOG.md to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

  • New feature (non-breaking change which adds functionality)
  • Optimization (back-end change that speeds up the code)
  • Bug fix (non-breaking change which fixes an issue)

Key checklist:

  • No style issues: $ flake8
  • All tests pass: $ python run-tests.py --unit
  • The documentation builds: $ cd docs and then $ make clean; make html

You can run all three at once, using $ python run-tests.py --quick.

Further checks:

  • Code is commented, particularly in hard-to-understand areas
  • Tests added that prove fix is effective or that feature works

@codecov
Copy link

codecov bot commented Dec 27, 2021 •
edited
Loading

Codecov Report

Merging #1874 (ceb47d0) into develop (ccf5be7) will increase coverage by 0.02%.
The diff coverage is 100.00%.

Impacted file tree graph

@@             Coverage Diff             @@
##           develop    #1874      +/-   ##
===========================================
+ Coverage    99.26%   99.28%   +0.02%     
===========================================
  Files          345      345              
  Lines        19107    19106       -1     
===========================================
+ Hits         18966    18970       +4     
+ Misses         141      136       -5
Impacted Files Coverage Δ
pybamm/__init__.py 94.95% <100.00%> (ø)
pybamm/util.py 100.00% <100.00%> (+2.87%) ⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update ccf5be7...ceb47d0. Read the comment docs.

brosaplanella
brosaplanella approved these changes Dec 27, 2021
View reviewed changes
Copy link
Member

@brosaplanella brosaplanella left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good, thanks Saransh! For a more robust fix, would it be possible to check if the jax version is the correct one and if not, throw a warning and treat the code as if jax wasn't installed?

@Saransh-cpp
Copy link
Member Author

Looks good, thanks Saransh! For a more robust fix, would it be possible to check if the jax version is the correct one and if not, throw a warning and treat the code as if jax wasn't installed?

I was trying to uninstall the existing version of jax while installing pybamm through a function but it felt weird editing a user's environment. The warning sounds really good! I'll try adding that instead of editing the notebooks.

@Saransh-cpp Saransh-cpp marked this pull request as draft December 27, 2021 20:44
@Saransh-cpp Saransh-cpp changed the title Make Colab notebooks work Make PyBaMM importable in Linux systems if jax is already installed Dec 28, 2021
@Saransh-cpp Saransh-cpp marked this pull request as ready for review December 28, 2021 08:41
Saransh-cpp
Saransh-cpp commented Dec 28, 2021
View reviewed changes
pybamm/util.py
@@ -356,6 +367,11 @@ def install_jax():

if system() == "Windows":
raise NotImplementedError("Jax is not available on Windows")
elif importlib.util.find_spec("jax") is not None:
if not is_jax_compatible():
raise ValueError(
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Codecov shows that this complete function is not included in the coverage, should I add # pragma: no cover for this function?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That sounds good

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, this is now an error. Should we make it a warning and proceed without jax?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PyBaMM installation (and importing PyBaMM) proceeds without jax if the installed jax version is not compatible. But this function throws an error if a user tries to install jax and has an incompatible jax version already installed.

Should it instead uninstall the existing version and install the compatible version?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I misunderstood! That works as it is :)

@priyanshuone6
Copy link
Member

priyanshuone6 commented Dec 28, 2021 •
edited
Loading

We should create a global variable for jax version, jaxlib version and use it in install_jax, is_jax_compatible

@brosaplanella
Copy link
Member

Looks great now, I will wait until tomorrow in case someone else wants to give it a look.

@brosaplanella
Copy link
Member

@all-contributors please add @Saransh-cpp for maintenance

@allcontributors allcontributors bot mentioned this pull request Dec 28, 2021
Merged
@allcontributors
Copy link
Contributor

@brosaplanella

I've put up a pull request to add @Saransh-cpp! 🎉

@brosaplanella brosaplanella merged commit 660c5d2 into pybamm-team:develop Dec 29, 2021
@Saransh-cpp Saransh-cpp deleted the make-colab-books-work branch December 29, 2021 17:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@brosaplanella brosaplanella brosaplanella approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@Saransh-cpp @priyanshuone6 @brosaplanella